#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include "tree.h"

static node * _alloc_node(int data)
{
	node *p = (node *)malloc(sizeof(struct node));
	if(!p){
		fprintf(stderr,"malloc failed !!!\n");
		exit(9);
	}
	p->left = p->right = p->parent = NULL;
	p->bf = 0;
	p->data = data;
	return p;
}

node * avl_search(node *root, int data)
{
	if(!root){
		return NULL;
	}
	node *p = root;
	while(p){
		if(p->data > data){
			p = p->left;
		}else if(p->data < data){
			p = p->right;
		}else{
			return p;
		}
	}
	return NULL;
}

static node * l_rotate(node **rootp, node *p)
{
	assert(rootp && *rootp && p && p->right);
	node *q = p->right;
	p->right = q->left;
	q->left = p;
	if(!p->parent){
		*rootp = q;
	}else if(p->parent->left == p){
		p->parent->left = q;
	}else{
		p->parent->right = q;
	}
	q->parent = p->parent;
	p->parent = q;
	p->bf -= q->bf<0?q->bf:0;
	p->bf ++;
	q->bf += p->bf>0?p->bf:0;
	q->bf ++;
	return q;
}
static node * r_rotate(node **rootp, node *p)
{
	assert(rootp &&  *rootp && p && p->left);
	node *q = p->left;
	p->left = q->right;
	q->right = p;
	if(!p->parent){
		*rootp = q;
	}else if(p == p->parent->left){
		p->parent->left = q;
	}else{
		p->parent->right = q;
	}
	q->parent = p->parent;
	p->parent = q;
	p->bf -= q->bf>0?q->bf:0;
	p->bf --;
	q->bf += p->bf<0?p->bf:0;
	q->bf --;
	return q;
}

static node * l_balence(node **rootp, node *p)
{
	assert(rootp && *rootp && p && p->right && p->bf==-2);
	if(p->right->bf == 1){
		r_rotate(rootp, p->right);
	}
	return l_rotate(rootp, p);
}
static node * r_balence(node **rootp, node *p)
{
	assert(rootp && *rootp && p && p->left && p->bf==2);
	if(p->left->bf == -1){
		l_rotate(rootp, p->left);
	}
	return r_rotate(rootp, p);
}

static int insert_fixup(node **rootp, node *p)
{
	assert(rootp && *rootp && p);
	node *q = NULL;
	while((q = p->parent)){
		q->bf += (p==q->left?1:-1);
		if(q->bf == -2){
			q = l_balence(rootp, q);
		}else if(q->bf == 2){
			q = r_balence(rootp, q);
		}
		if(!q->bf){
			break;
		}
		p = q;
	}
	return 0;
}

node * avl_insert(node **rootp, int data)
{
	assert(rootp);

	if(!*rootp){
		*rootp = _alloc_node(data);
		return *rootp;
	}
	
	node *p = *rootp;
	node *q = NULL;
	while(p){
		q = p;
		if(p->data > data){
			p = p->left;
		}else if(p->data < data){
			p = p->right;
		}else{
			return p;
		}
	}
	p = _alloc_node(data);
	p->parent = q;
	if(q->data > data){
		q->left = p;
	}else{
		q->right = p;
	}
	insert_fixup(rootp, p);
	return p;
}

static int delete_fixup(node **rootp, node *p)
{
	assert(rootp && *rootp && p);
	node *q = NULL;
	while((q = p->parent)){
		q->bf -= (p==q->left?1:-1);
		if(q->bf == -2){
			q = l_rotate(rootp, q);
		}else if(q->bf == 2){
			q = r_rotate(rootp, q);
		}
		if(q->bf){
			break;
		}
		p = q;
	}
	return 0;
}

node * avl_delete(node **rootp, int data)
{
	node *p = avl_search(*rootp, data);
	if(!p){
		return NULL;
	}
	node *q = NULL;
	if(p->left && p->right){
		q = successor(p);
		if(q->right){
			q->right->parent = q->parent;
		}
		p->data = q->data;
		if(q == q->parent->left){
			q->parent->left = q->right;
		}else{
			q->parent->right = q->right;
		}
		delete_fixup(rootp, q);
		q->parent = NULL;
		return q;
	}

	if(p->right){
		q = p->right;
	}else{
		q = p->left;
	}
	if(q){
		q->parent = p->parent;
	}
	if(!p->parent){
		*rootp = q;
	}else if(p == p->parent->left){
		p->parent->left = q;
	}else{
		p->parent->right = q;
	}
	delete_fixup(rootp, p);
	p->parent = NULL;
	return p;
}
